import cv2
import mediapipe as mp
import math


class PoseMoudle():
    def __init__(self,
                 static_image_mode=False,
                 model_complexity=1,
                 smooth_landmarks=True,
                 enable_segmentation=False,
                 smooth_segmentation=True,
                 min_detection_confidence=0.5,
                 min_tracking_confidence=0.5):
        self.static_image_mode = static_image_mode # 是否是静态图片
        self.model_complexity = model_complexity # 1 或 2
        self.smooth_landmarks = smooth_landmarks # 是否平滑关键点
        self.enable_segmentation = enable_segmentation # 是否人体分割
        self.smooth_segmentation = smooth_segmentation # 是否平滑分割
        self.min_detection_confidence = min_detection_confidence # 置信度
        self.min_tracking_confidence = min_tracking_confidence # 跟踪置信度阈值

        self.mpPose = mp.solutions.pose
        self.pose = self.mpPose.Pose(self.static_image_mode, self.smooth_landmarks,
                                     self.enable_segmentation, self.enable_segmentation,
                                     self.smooth_segmentation, self.min_detection_confidence,
                                     self.min_tracking_confidence)
        self.mpDraw = mp.solutions.drawing_utils



    def findPose(self, frame, draw=True):
        """
        寻找姿势并画出
        :param frame:
        :param draw:
        :return:
        """
        frameRGB = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
        self.result = self.pose.process(frameRGB)
        if self.result.pose_landmarks:
            if draw:
                self.mpDraw.draw_landmarks(frame, self.result.pose_landmarks,
                                                  self.mpPose.POSE_CONNECTIONS)
        return frame


    def findPosition(self, frame, draw=True):
        """
         获取33个姿势点，并添加到listlm列表中
        :param frame:
        :param draw:
        :return:
        """
        self.listlm = []
        self.result = self.pose.process(frame)
        if self.result.pose_landmarks:
            for id, lm in enumerate(self.result.pose_landmarks.landmark):
                h,w,c = frame.shape
                cx, cy = int(lm.x*w), int(lm.y*h)
                self.listlm.append([id, cx, cy])
                if draw:
                    self.mpDraw.draw_landmarks(frame, self.result.pose_landmarks,self.mpPose.POSE_CONNECTIONS)
        return self.listlm


    def findAngle(self,frame,p1,p2,p3,draw=True):
        x1,y1 = self.listlm[p1][1:]
        x2, y2 = self.listlm[p2][1:]
        x3, y3 = self.listlm[p3][1:]
        # 计算角度
        angle = math.degrees(math.atan2(y3-y2,x3-x2) - math.atan2(y1-y2,x1-x2))
        if angle<0:
            angle += 360
        if draw:
            cv2.line(frame,(x1,y1),(x2,y2),(255,255,255),3)
            cv2.line(frame, (x3, y3), (x2, y2), (255, 255, 255), 3)
            cv2.circle(frame, (x1,y1),10,(0,0,255),cv2.FILLED)
            cv2.circle(frame, (x1, y1), 15, (0, 0, 255), 3)
            cv2.circle(frame, (x2, y2), 10, (0, 0, 255), cv2.FILLED)
            cv2.circle(frame, (x2, y2), 15, (0, 0, 255), 3)
            cv2.circle(frame, (x3, y3), 10, (0, 0, 255), cv2.FILLED)
            cv2.circle(frame, (x3, y3), 15, (0, 0, 255), 3)
            cv2.putText(frame, str(int(angle)),(x2-50,y2+50),cv2.FONT_HERSHEY_PLAIN,2,(0,255,0),2)
        return angle